"""
抖音视频下载器 - 高级版
针对抖音复杂的视频加载机制进行优化
"""

import os
import time
import requests
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.chrome import ChromeDriverManager

def download_douyin_video(url, headless=False):
    """
    使用 Selenium 下载抖音视频 - 高级版本
    
    参数:
        url (str): 抖音视频的分享链接
        headless (bool): 是否使用无头模式（不显示浏览器窗口）
    """
    print("=" * 50)
    print(f"开始处理抖音视频: {url}")
    
    # 设置 Chrome 选项
    chrome_options = Options()
    if headless:
        chrome_options.add_argument("--headless")  # 无头模式
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    chrome_options.add_argument("--disable-infobars")
    chrome_options.add_argument("--disable-extensions")
    chrome_options.add_argument("--disable-notifications")
    chrome_options.add_argument("--disable-software-rasterizer")
    chrome_options.add_argument("--disable-webgl")
    chrome_options.add_experimental_option('excludeSwitches', ['enable-logging'])
    
    # 设置用户代理
    user_agent = "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1"
    chrome_options.add_argument(f"user-agent={user_agent}")
    
    # 设置窗口大小（模拟移动设备）
    chrome_options.add_argument("--window-size=375,812")
    
    driver = None
    try:
        print("初始化浏览器驱动...")
        # 使用 webdriver-manager 自动管理驱动
        service = Service(ChromeDriverManager().install())
        driver = webdriver.Chrome(service=service, options=chrome_options)
        
        print("正在访问页面...")
        driver.get(url)
        
        # 等待页面加载
        print("等待页面加载...")
        time.sleep(3)
        
        # 检查是否需要登录
        if "登录" in driver.page_source:
            print("检测到登录页面，需要手动登录...")
            if headless:
                print("无头模式下无法登录，请使用非无头模式(headless=False)")
                return None
            else:
                print("请在浏览器中完成登录，完成后按回车继续...")
                input()
        
        # 获取页面标题
        title = driver.title.replace(" - 抖音", "").strip()
        print(f"视频标题: {title}")
        
        # 尝试点击播放按钮（如果存在）
        try:
            play_button = WebDriverWait(driver, 10).until(
                EC.element_to_be_clickable((By.CSS_SELECTOR, 'div[data-e2e="video-player"] button'))
            )
            play_button.click()
            print("点击了播放按钮")
            time.sleep(2)
        except:
            print("未找到播放按钮或无需点击")
        
        # 查找视频元素 - 使用更可靠的方法
        print("查找视频元素...")
        
        # 方法1: 使用JavaScript获取视频元素
        video_elements = driver.execute_script("""
            // 尝试获取所有video元素
            let videos = document.querySelectorAll('video');
            if (videos.length > 0) return videos;
            
            // 尝试在Shadow DOM中查找
            let shadowRoots = [];
            document.querySelectorAll('*').forEach(el => {
                if (el.shadowRoot) shadowRoots.push(el.shadowRoot);
            });
            
            for (let root of shadowRoots) {
                videos = root.querySelectorAll('video');
                if (videos.length > 0) return videos;
            }
            
            // 尝试在iframe中查找
            let iframes = document.querySelectorAll('iframe');
            for (let iframe of iframes) {
                try {
                    let doc = iframe.contentDocument || iframe.contentWindow.document;
                    videos = doc.querySelectorAll('video');
                    if (videos.length > 0) return videos;
                    
                    // 在iframe的Shadow DOM中查找
                    doc.querySelectorAll('*').forEach(el => {
                        if (el.shadowRoot) shadowRoots.push(el.shadowRoot);
                    });
                    
                    for (let root of shadowRoots) {
                        videos = root.querySelectorAll('video');
                        if (videos.length > 0) return videos;
                    }
                } catch (e) {}
            }
            
            return [];
        """)
        
        # 方法2: 如果JS方法失败，使用传统方法
        if not video_elements:
            print("使用传统方法查找视频元素...")
            video_elements = driver.find_elements(By.TAG_NAME, 'video')
        
        if not video_elements:
            print("无法找到视频元素，保存页面截图...")
            driver.save_screenshot("debug.png")
            print("请查看 debug.png 文件进行故障排除")
            return None
        
        print(f"找到 {len(video_elements)} 个视频元素")
        
        # 获取视频URL
        video_url = None
        for i, video in enumerate(video_elements):
            # 使用JavaScript获取视频源
            src = driver.execute_script("""
                return arguments[0].src || 
                       arguments[0].getAttribute('data-src') || 
                       arguments[0].getAttribute('srcset') || 
                       arguments[0].querySelector('source')?.src;
            """, video)
            
            print(f"视频元素 {i+1} 的URL: {src}")
            
            if src and "http" in src:
                video_url = src
                print(f"使用视频元素 {i+1}")
                break
        
        if not video_url:
            print("无法获取有效的视频URL")
            return None
        
        print(f"最终视频URL: {video_url}")
        
        # 下载视频
        headers = {
            "User-Agent": user_agent,
            "Referer": driver.current_url,
            "Accept": "*/*",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Range": "bytes=0-"
        }
        
        print("开始下载视频...")
        start_time = time.time()
        
        # 使用requests下载视频
        response = requests.get(video_url, headers=headers, stream=True)
        
        if response.status_code == 200 or response.status_code == 206:
            # 创建下载目录
            os.makedirs("downloads", exist_ok=True)
            
            # 生成安全文件名
            safe_title = "".join(c for c in title if c.isalnum() or c in " -_")
            if not safe_title:
                safe_title = "douyin_video"
            
            filename = f"downloads/{safe_title}_{int(time.time())}.mp4"
            
            # 获取文件大小
            file_size = int(response.headers.get('Content-Length', 0))
            
            # 下载进度跟踪
            downloaded = 0
            chunk_size = 1024 * 1024  # 1MB
            
            with open(filename, 'wb') as f:
                for chunk in response.iter_content(chunk_size=chunk_size):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if file_size > 0:
                            percent = (downloaded / file_size) * 100
                            print(f"下载进度: {percent:.1f}%", end='\r')
            
            download_time = time.time() - start_time
            actual_size = os.path.getsize(filename) / (1024 * 1024)  # MB
            
            print(f"\n视频已保存至: {filename}")
            print(f"文件大小: {actual_size:.2f} MB")
            print(f"下载时间: {download_time:.2f} 秒")
            return filename
        else:
            print(f"下载失败，状态码: {response.status_code}")
            print(f"响应头: {response.headers}")
            return None
    except Exception as e:
        print(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        if driver:
            driver.quit()
            print("浏览器已关闭")

if __name__ == "__main__":
    # 抖音视频链接
    share_url = "https://v.douyin.com/X01lRooQLDc/"
    
    # 使用非无头模式以便调试
    result = download_douyin_video(share_url, headless=False)
    
    if result:
        print(f"成功下载视频: {result}")
    else:
        print("视频下载失败")